package com.ruoyix.novel.controller;


import cn.dev33.satoken.annotation.SaIgnore;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

@RestController
@CrossOrigin(origins = "*")
public class ChatListController {

    private final WebClient webClient;
    private final ObjectMapper objectMapper = new ObjectMapper();

    public ChatListController() {
        this.webClient = WebClient.builder()
            .baseUrl("https://ark.cn-beijing.volces.com")
            .defaultHeader("Authorization", "Bearer f9c1148b-26fc-4c92-8977-fca871de1af3")
            .build();
    }

    @SaIgnore
    @PostMapping(value = "/api1/chat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@RequestBody ChatRequest request) {
        return Flux.fromIterable(request.getQuestions())
            .index()
            .concatMap(tuple -> {
                long index = tuple.getT1() + 1;
                String question = tuple.getT2();

                // 生成问题事件
                String questionJson = String.format("{\"question_index\": %d, \"question\": %s}",
                    index, toJsonString(question));
                String questionEvent = "data: " + questionJson + "\n\n";

                // 调用外部API
                Flux<String> apiResponse = callExternalApi(question, index);

                // 处理API响应
                List<String> currentAnswer = Collections.synchronizedList(new ArrayList<>());
                Flux<String> processedResponse = apiResponse
                    .doOnNext(chunk -> processChunk(chunk, currentAnswer))
                    .doOnComplete(() -> saveAnswer(currentAnswer, question, index));

                return Flux.just(questionEvent).concatWith(processedResponse);
            })
            .concatWithValues("data: [DONE]\n\n");
    }

    private Flux<String> callExternalApi(String question, long index) {
        return webClient.post()
            .uri("/api/v3/chat/completions")
            .contentType(MediaType.APPLICATION_JSON)
            .bodyValue(createRequestBody(question))
            .retrieve()
            .onStatus(status -> !status.is2xxSuccessful(),
                response -> response.bodyToMono(String.class)
                    .flatMap(error -> Mono.error(new RuntimeException("API Error: " + error))))
            .bodyToFlux(String.class)
            .filter(text -> !text.trim().isEmpty())
            .onErrorResume(e -> Flux.just(
                String.format("data: {\"error\": \"%s\", \"question_index\": %d}\n\n",
                    e.getMessage(), index)));
    }

    private Map<String, Object> createRequestBody(String question) {
        return Map.of(
            "model", "doubao-1-5-pro-256k-250115",
            "stream", true,
            "messages", List.of(
                Map.of("role", "system", "content", "You are a helpful assistant."),
                Map.of("role", "user", "content", question)
            )
        );
    }

    private void processChunk(String chunk, List<String> currentAnswer) {
        try {
            if (chunk.startsWith("data: ")) {
                String jsonStr = chunk.substring(5).trim();
                if (jsonStr.equals("[DONE]")) return;

                JsonNode node = objectMapper.readTree(jsonStr);
                JsonNode choices = node.get("choices");
                if (choices != null && choices.isArray() && choices.size() > 0) {
                    JsonNode delta = choices.get(0).get("delta");
                    if (delta != null) {
                        JsonNode content = delta.get("content");
                        if (content != null && content.isTextual()) {
                            currentAnswer.add(content.asText());
                        }
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private void saveAnswer(List<String> answerParts, String question, long index) {
        String answer = String.join("", answerParts);
        String content = String.format("问题%d: %s\n答案：%s\n\n", index, question, answer);

        Mono.fromRunnable(() -> {
                try {
                    Files.write(Paths.get("answers.txt"),
                        content.getBytes(StandardCharsets.UTF_8),
                        StandardOpenOption.CREATE,
                        StandardOpenOption.APPEND);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            })
            .subscribeOn(Schedulers.boundedElastic())
            .subscribe();
    }

    private String toJsonString(String str) {
        try {
            return objectMapper.writeValueAsString(str);
        } catch (Exception e) {
            return "\"\"";
        }
    }

    @Data
    static class ChatRequest {
        private List<String> questions;
    }

    @SaIgnore
    @PostMapping(value = "/api1/chat2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat2(@RequestBody ChatRequest request) {
        // 创建一个包含所有问题处理的Flux列表
        List<Flux<String>> questionFluxes = new ArrayList<>();
        int questionCount = request.getQuestions().size();

        // 为每个问题创建处理流程
        for (int i = 0; i < questionCount; i++) {
            final int index = i + 1;
            final String question = request.getQuestions().get(i);

            // 创建单个问题的处理流程
            Flux<String> singleQuestionFlow = createQuestionFlow(question, index);
            questionFluxes.add(singleQuestionFlow);
        }

        // 按顺序拼接所有问题的处理流程，最后添加结束标记
        return Flux.concat(questionFluxes)
            .concatWithValues("data: [DONE]\n\n");
    }

    private Flux<String> createQuestionFlow(String question, int index) {
        // 生成问题事件
        String questionJson = String.format("{\"question_index\": %d, \"question\": %s}",
            index, toJsonString(question));
        String questionEvent = "data: " + questionJson + "\n\n";

        // 调用外部API并处理响应
        return Flux.just(questionEvent)
            .concatWith(processQuestionAnswer(question, index));
    }

    private Flux<String> processQuestionAnswer(String question, int index) {
        // 调用API并处理响应
        List<String> currentAnswer = Collections.synchronizedList(new ArrayList<>());

        return callExternalApi(question, index)

            .doOnNext(chunk -> processChunk(chunk, currentAnswer))
            .doOnComplete(() -> saveAnswer(currentAnswer, question, index))
            .onErrorResume(e -> Flux.just(
                String.format("data: {\"error\": \"%s\", \"question_index\": %d}\n\n",
                    e.getMessage(), index)));
    }
}
